#include "ohos_video_h264_decoder.h"

#include "../../videofilters/h26x/h264-utils.h"
#include "bctoolbox/crypto.h"
#include "mediastreamer2/msticker.h"

using namespace ohos;

OHOSVideoH264Decoder::OHOSVideoH264Decoder(MSFilter *f)
    : mFilter(f), mUnpacker(0), mSPS(0), mPPS(0), mYUVMsg(0), mBitstream(0), mBitstreamSize(65536),
      mLastErrorReportTime(0), mWidth(MS_VIDEO_SIZE_UNKNOWN_W), mHeight(MS_VIDEO_SIZE_UNKNOWN_H), mInitialized(false),
      mFirstImageDecoded(false) {
}

OHOSVideoH264Decoder::~OHOSVideoH264Decoder() {
	if (mBitstream != 0) {
		ms_free(mBitstream);
	}
	uninitialize();
}

void OHOSVideoH264Decoder::initialize() {
	if (!mInitialized) {
		mBitstream = static_cast<uint8_t *>(ms_malloc0(mBitstreamSize));

		mFirstImageDecoded = false;
		mUnpacker = rfc3984_new_with_factory(mFilter->factory);

		data_ = std::make_shared<ohos::CodecData>();
		format_ = std::make_shared<ohos::VideoFormat>();

		format_->width = 1920;
		format_->height = 1920;
		format_->codecMime = OH_AVCODEC_MIMETYPE_VIDEO_AVC;
		format_->pixelFormat = AV_PIXEL_FORMAT_NV21;
		format_->lowLatency = true;

		decoder_ = OH_VideoDecoder_CreateByMime(format_->codecMime.c_str());
		if (decoder_ == NULL) {
			ms_message("create video decoder failed");
			return;
		}

		OH_AVFormat *avformat = OH_AVFormat_Create();
		OH_AVFormat_SetStringValue(avformat, OH_MD_KEY_CODEC_MIME, format_->codecMime.c_str());
		OH_AVFormat_SetIntValue(avformat, OH_MD_KEY_WIDTH, format_->width);
		OH_AVFormat_SetIntValue(avformat, OH_MD_KEY_HEIGHT, format_->height);
		OH_AVFormat_SetIntValue(avformat, OH_MD_KEY_PIXEL_FORMAT, format_->pixelFormat);

		int32_t ret = OH_VideoDecoder_Configure(decoder_, avformat);
		if (ret != AV_ERR_OK) {
			ms_message("set configure failed");
			return;
		}
		OH_AVFormat_Destroy(avformat);
		avformat = nullptr;

		OH_AVCodecCallback cb = {&DataCallback::OnCodecError, &DataCallback::OnCodecFormatChange,
		                         &DataCallback::OnNeedInputBuffer, &DataCallback::OnNewOutputBuffer};
		ret = OH_VideoDecoder_RegisterCallback(decoder_, cb, this);
		if (ret != AV_ERR_OK) {
			ms_message("register callback failed");
			return;
		}

		ret = OH_VideoDecoder_Prepare(decoder_);
		if (ret != AV_ERR_OK) {
			ms_message("prepare failed");
			return;
		}

		if (decoder_ == nullptr) {
			return;
		}
		if (OH_VideoDecoder_Start(decoder_) != AV_ERR_OK) {
		}
		ms_average_fps_init(&mFPS, "OpenH264 decoder: FPS=%f");
		mInitialized = true;
		ms_message("ohos video decoder init success");
	}
}

void OHOSVideoH264Decoder::feedNV21AsI420(uint8_t *addr, int width, int height, int stride) {
	if (!mYUVMsg || width != mWidth || height != mHeight) {
		if (mYUVMsg) freemsg(mYUVMsg);
		mWidth = width;
		mHeight = height;
		mYUVMsg = ms_yuv_buf_alloc(&mOutbuf, width, height);
	}

	uint8_t *yPlane = mOutbuf.planes[0];
	uint8_t *uPlane = mOutbuf.planes[1];
	uint8_t *vPlane = mOutbuf.planes[2];

	int yStrideDst = mOutbuf.strides[0];
	int uvStrideDst = mOutbuf.strides[1];

	int ySize = stride_ * height;
	uint8_t *uvSrc = addr + ySize;

	for (int j = 0; j < height; j++) {
		memcpy(yPlane + j * yStrideDst, addr + j * stride_, width);
	}

	for (int j = 0; j < height / 2; j++) {
		uint8_t *srcRow = uvSrc + j * stride_;
		uint8_t *uDstRow = uPlane + j * uvStrideDst;
		uint8_t *vDstRow = vPlane + j * uvStrideDst;

		for (int i = 0; i < width / 2; i++) {
			vDstRow[i] = srcRow[2 * i];
			uDstRow[i] = srcRow[2 * i + 1];
		}
	}

	ms_queue_put(mFilter->outputs[0], dupmsg(mYUVMsg));
}

void OHOSVideoH264Decoder::feed() {
	if (!isInitialized()) {
		ms_error("OHOSVideoH264Decoder::feed(): not initialized");
		ms_queue_flush(mFilter->inputs[0]);
		return;
	}

	MSQueue nalus;
	ms_queue_init(&nalus);

	mblk_t *im;
	bool requestPLI = false;
	while ((im = ms_queue_get(mFilter->inputs[0])) != NULL) {
		unsigned int ret;
		if ((mSPS != 0) && (mPPS != 0)) {
			rfc3984_unpack_out_of_band_sps_pps(mUnpacker, mSPS, mPPS);
			mSPS = NULL;
			mPPS = NULL;
		}
		ret = rfc3984_unpack2(mUnpacker, im, &nalus);
		if (ret & Rfc3984FrameAvailable) {
			int len = nalusToFrame(&nalus);

			if (ret & Rfc3984FrameCorrupted) requestPLI = true;

			OH_AVCodecBufferAttr attr = {0, 0, 0, AVCODEC_BUFFER_FLAGS_NONE};
			attr.pts = im->timestamp.tv_sec * 1000000 + im->timestamp.tv_usec;
			attr.size = len;
			attr.flags = 0;
			int indx = -1;
			int ret = data_->InputData(mBitstream, len, &attr, &indx, std::chrono::milliseconds(8));
			if (ret != ERRCODE_OK) {
				OH_LOG_Print(LOG_APP, LOG_ERROR, LOG_DOMAIN, "VideoDecoder", "codec data set InputData failed,\
                        errcode=%{public}d",
				             ret);
				break;
			}
			FreeInputData(indx);

			if (GetFormatInfo(ohos::FormatType::FORMAT_TYPE_IMAGE_WIDTH) != mWidth ||
			    GetFormatInfo(ohos::FormatType::FORMAT_TYPE_IMAGE_HEIGHT) != mHeight) {
				mWidth = GetFormatInfo(ohos::FormatType::FORMAT_TYPE_IMAGE_WIDTH);
				mHeight = GetFormatInfo(ohos::FormatType::FORMAT_TYPE_IMAGE_HEIGHT);
				stride_ = GetFormatInfo(ohos::FormatType::FORMAT_TYPE_VIDEO_STRIDE);

				if (mYUVMsg) {
					freemsg(mYUVMsg);
				}
				mYUVMsg = ms_yuv_buf_alloc(&mOutbuf, mWidth, mHeight);
				ms_filter_notify_no_arg(mFilter, MS_FILTER_OUTPUT_FMT_CHANGED);
				ms_message("%s: %d mWidth: %d, mHeight: %d, stride_: %d", __FILE__, __LINE__, mWidth, mHeight, stride_);
			}

			ohos::AvBufferInfo info{0, nullptr};
			ret = data_->OutputData(info, std::chrono::milliseconds(8));
			if (ret != ohos::ERRCODE_OK) {
				break;
			}

			const OH_AVCodecBufferAttr *attr1 = info.Attr();
			if (attr1 == nullptr) {
				ms_message("output data is null");
				FreeOutputData(info.Index());
				break;
			}

			uint8_t *addr = info.AvBufferAddr();
			if (addr == nullptr) {
				ms_message("output data addr is null");
				FreeOutputData(info.Index());
				break;
			}

			feedNV21AsI420(addr, mWidth, mHeight, stride_);

			if (ms_average_fps_activity(&mFPS, mFilter->ticker->time, TRUE)) {
				ms_message("OpenH264 decoder: Frame size: %dx%d", mWidth, mHeight);
			}

			if (!mFirstImageDecoded) {
				mFirstImageDecoded = true;
				ms_filter_notify_no_arg(mFilter, MS_VIDEO_DECODER_FIRST_IMAGE_DECODED);
			}

			FreeOutputData(info.Index());
		}
	}
	if (requestPLI) {
		if (mAVPFEnabled) {
			ms_filter_notify_no_arg(mFilter, MS_VIDEO_DECODER_SEND_PLI);
		} else if (((mFilter->ticker->time - mLastErrorReportTime) > 5000) || (mLastErrorReportTime == 0)) {
			mLastErrorReportTime = mFilter->ticker->time;
			ms_filter_notify_no_arg(mFilter, MS_VIDEO_DECODER_DECODING_ERRORS);
		}
	}
}

void OHOSVideoH264Decoder::uninitialize() {

    if (avFormat_ != nullptr) {
        OH_AVFormat_Destroy(avFormat_);
        avFormat_ = nullptr;
    }

    if (decoder_ != nullptr) {
        OH_VideoDecoder_Destroy(decoder_);
        decoder_ = nullptr;
    }

	if (mSPS != 0) {
		freemsg(mSPS);
		mSPS = NULL;
	}
	if (mPPS != 0) {
		freemsg(mPPS);
		mPPS = NULL;
	}
	if (mYUVMsg != 0) {
		freemsg(mYUVMsg);
		mYUVMsg = NULL;
	}

	if (mUnpacker) {
		rfc3984_destroy(mUnpacker);
		mUnpacker = NULL;
	}
    
	mInitialized = false;
}

void OHOSVideoH264Decoder::provideSpropParameterSets(char *value, size_t valueSize) {
	char *b64_sps = value;
	char *b64_pps = strchr(value, ',');
	if (b64_pps) {
		*b64_pps = '\0';
		++b64_pps;
		ms_message("OpenH264 decoder: Got sprop-parameter-sets sps=%s, pps=%s", b64_sps, b64_pps);
		mSPS = allocb(valueSize, 0);
		bctbx_base64_decode(mSPS->b_wptr, &valueSize, (unsigned char *)b64_sps, strlen(b64_sps));
		mSPS->b_wptr += valueSize;
		mPPS = allocb(valueSize, 0);
		bctbx_base64_decode(mPPS->b_wptr, &valueSize, (unsigned char *)b64_pps, strlen(b64_pps));
		mPPS->b_wptr += valueSize;
	}
}

void OHOSVideoH264Decoder::resetFirstImageDecoded() {
	mFirstImageDecoded = false;
	mWidth = MS_VIDEO_SIZE_UNKNOWN_W;
	mHeight = MS_VIDEO_SIZE_UNKNOWN_H;
}

MSVideoSize OHOSVideoH264Decoder::getSize() const {
	MSVideoSize size;
	size.width = mWidth;
	size.height = mHeight;
	return size;
}

float OHOSVideoH264Decoder::getFps() const {
	return ms_average_fps_get(&mFPS);
}

const MSFmtDescriptor *OHOSVideoH264Decoder::getOutFmt() const {
	MSVideoSize vsize = {mWidth, mHeight};
	return ms_factory_get_video_format(mFilter->factory, "YUV420P", vsize, 0, NULL);
}

int OHOSVideoH264Decoder::nalusToFrame(MSQueue *nalus) {
	mblk_t *im;
	uint8_t *dst = mBitstream;
	uint8_t *end = mBitstream + mBitstreamSize;
	bool startPicture = true;

	while ((im = ms_queue_get(nalus)) != NULL) {
		uint8_t *src = im->b_rptr;
		int nalLen = im->b_wptr - src;
		if ((dst + nalLen + 128) > end) {
			int pos = dst - mBitstream;
			enlargeBitstream(mBitstreamSize + nalLen + 128);
			dst = mBitstream + pos;
			end = mBitstream + mBitstreamSize;
		}
		if ((src[0] == 0) && (src[1] == 0) && (src[2] == 0) && (src[3] == 1)) {
			// Workaround for stupid RTP H264 sender that includes nal markers
#if MSOPENH264_DEBUG
			ms_warning("OpenH264 decoder: stupid RTP H264 encoder");
#endif
			int size = im->b_wptr - src;
			memcpy(dst, src, size);
			dst += size;
		} else {
			uint8_t naluType = *src & 0x1f;
#if MSOPENH264_DEBUG
			if ((naluType != 1) && (naluType != 7) && (naluType != 8)) {
				ms_message("OpenH264 decoder: naluType=%d", naluType);
			}
			if (naluType == 7) {
				ms_message("OpenH264 decoder: Got SPS");
			}
			if (naluType == 8) {
				ms_message("OpenH264 decoder: Got PPS");
			}
#endif
			if (startPicture || (naluType == 6)              // SEI
			    || (naluType == 7)                           // SPS
			    || (naluType == 8)                           // PPS
			    || ((naluType >= 14) && (naluType <= 18))) { // Reserved
				*dst++ = 0;
				startPicture = false;
			}

			// Prepend nal marker
			*dst++ = 0;
			*dst++ = 0;
			*dst++ = 1;
			*dst++ = *src++;
			while (src < (im->b_wptr - 3)) {
				if ((src[0] == 0) && (src[1] == 0) && (src[2] < 3)) {
					*dst++ = 0;
					*dst++ = 0;
					*dst++ = 3;
					src += 2;
				}
				*dst++ = *src++;
			}
			while (src < im->b_wptr) {
				*dst++ = *src++;
			}
		}
		freemsg(im);
	}
	return dst - mBitstream;
}

void OHOSVideoH264Decoder::enlargeBitstream(int newSize) {
	mBitstreamSize = newSize;
	mBitstream = static_cast<uint8_t *>(ms_realloc(mBitstream, mBitstreamSize));
}

void OHOSVideoH264Decoder::DataCallback::OnCodecError(OH_AVCodec *codec, int32_t errorCode, void *userData) {
	(void)codec;
	OHOSVideoH264Decoder *decoder = static_cast<OHOSVideoH264Decoder *>(userData);
	decoder->errorCode_ = errorCode;
	OH_LOG_Print(LOG_APP, LOG_ERROR, LOG_DOMAIN, "Decoder", "On decoder error, error code: %{public}d", errorCode);
}

void OHOSVideoH264Decoder::DataCallback::OnCodecFormatChange(OH_AVCodec *codec, OH_AVFormat *format, void *userData) {
	(void)codec;
	if (userData == nullptr || format == nullptr) {
		OH_LOG_Print(LOG_APP, LOG_ERROR, LOG_DOMAIN, "Decoder", "OnCodecFormatChange error, userData/format is NULL");
		return;
	}

	OHOSVideoH264Decoder *decoder = static_cast<OHOSVideoH264Decoder *>(userData);
	if (decoder->avFormat_ == nullptr) {
		decoder->avFormat_ = OH_AVFormat_Create();
	}
	OH_AVFormat_Copy(decoder->avFormat_, format);
	decoder->fmtChanged_ = true;
}

void OHOSVideoH264Decoder::DataCallback::OnNeedInputBuffer(OH_AVCodec *codec,
                                                           uint32_t index,
                                                           OH_AVBuffer *buffer,
                                                           void *userData) {

	(void)codec;
	if (userData == nullptr || buffer == nullptr) {
		OH_LOG_Print(LOG_APP, LOG_ERROR, LOG_DOMAIN, "Decoder", "OnNeedInputBuffer error, userData is NULL");
		return;
	}
	OHOSVideoH264Decoder *decoder = static_cast<OHOSVideoH264Decoder *>(userData);
	CodecData *codecData = decoder->GetCodecData();
	if (codecData != nullptr) {
		codecData->NeedInputBuffer(index, buffer);
	}
}

CodecData *OHOSVideoH264Decoder::GetCodecData() {
	AVCODEC_CHECK(data_ == nullptr, "codecData is NULL", nullptr);
	return data_.get();
}

int32_t OHOSVideoH264Decoder::FreeInputData(int32_t index) {
	AVCODEC_CHECK(decoder_ == nullptr, "decoder is NULL", -1);
	int32_t ret = OH_VideoDecoder_PushInputBuffer(decoder_, index);
	if (ret != AV_ERR_OK) {
		errorCode_ = ret;
		return ERRCODE_UNKNOWN;
	}

	return ERRCODE_OK;
}

int32_t OHOSVideoH264Decoder::GetFormatInfo(FormatType type) {
	int result = 0;

	const char *const formatKeys[FORMAT_TYPE_NBR] = {
	    OH_MD_KEY_VIDEO_PIC_WIDTH,
	    OH_MD_KEY_VIDEO_PIC_HEIGHT,
	    OH_MD_KEY_VIDEO_STRIDE,
	    OH_MD_KEY_VIDEO_SLICE_HEIGHT,
	    OH_MD_KEY_PIXEL_FORMAT,
	    OH_MD_KEY_VIDEO_CROP_TOP,
	    OH_MD_KEY_VIDEO_CROP_BOTTOM,
	    OH_MD_KEY_VIDEO_CROP_LEFT,
	    OH_MD_KEY_VIDEO_CROP_RIGHT,
	    OH_MD_KEY_WIDTH,
	    OH_MD_KEY_HEIGHT,
	};

	AVCODEC_CHECK(avFormat_ == nullptr, "avFormat_ is not initted!", ERRCODE_UNKNOWN);
	AVCODEC_CHECK((type < 0 || type >= FORMAT_TYPE_NBR), "FormatType ERROR!", ERRCODE_PARAM_INVALID);

	bool ret = OH_AVFormat_GetIntValue(avFormat_, formatKeys[type], &result);

	return ret ? result : -1;
}

void OHOSVideoH264Decoder::DataCallback::OnNewOutputBuffer(OH_AVCodec *codec,
                                                           uint32_t index,
                                                           OH_AVBuffer *buffer,
                                                           void *userData) {
	(void)codec;
	if (userData == nullptr || buffer == nullptr) {
		OH_LOG_Print(LOG_APP, LOG_ERROR, LOG_DOMAIN, "Decoder",
		             "OnNewOutputBuffer error, userData/OH_AVBuffer is NULL");
		return;
	}
	OHOSVideoH264Decoder *decoder = static_cast<OHOSVideoH264Decoder *>(userData);
	CodecData *codecData = decoder->GetCodecData();
	if (codecData != nullptr) {
		codecData->NewOutputBuffer(index, buffer);
	}
}

int32_t OHOSVideoH264Decoder::FreeOutputData(uint32_t index) {
	AVCODEC_CHECK(decoder_ == nullptr, "decoder is NULL", -1);

	int ret = OH_VideoDecoder_FreeOutputBuffer(decoder_, index);
	if (ret != AV_ERR_OK) {
		OH_LOG_Print(LOG_APP, LOG_ERROR, LOG_DOMAIN, "VideoDecoder", "OH_VideoDecoder_FreeOutputBuffer failed");
		errorCode_ = ret;
		return -1;
	}

	return 0;
}
